//065959264	Limor Bagizada	040199770	Sarit Chicotay

package Base;

import java.util.ArrayList;
import java.util.List;
import java.util.Map.Entry;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.TreeSet;

/**
 * 
 * Corpus class represents a corpus and its attributes
 *
 */
public class Corpus {
	
	// n-gram
	private int nGram;
	
	// ArrayList of all events in the corpus (not distinct)
	private ArrayList<NGram> words;
	
	// map of the distinct events in the document and the number of appearances for each
	// bigram - number of appearances for each event, unigram - number of appearances for each word 
	private SortedMap<NGram, Integer> vocabulary;
	
	// map of the distinct words and the number of appearances for each word (required for calculating the corpus size in bigram)
	private SortedMap<String, Integer> vocabulary_unigram;

	// map: 
	// key - an integer (r) which denote a frequency
	// value - a set of all events which appear r times in the corpus
	private SortedMap<Integer, TreeSet<NGram>> count_r;
	
	// Constructor
	public Corpus(List<String> words)
	{
		this(words, 1);
	}
	
	/**
	 * Constructor
	 * @param words - list of words in the corpus
	 */
	public Corpus(List<String> content, int nGram)
	{
		// n-gram
		this.nGram = nGram;
		
		
		// initialization
		this.words = new ArrayList<NGram>();
		this.vocabulary = new TreeMap<NGram, Integer>();
		this.vocabulary_unigram = new TreeMap<String, Integer>();

		// remove empty words
		for (String word : content)
		{
			if (word.trim().equals(""))
				content.remove(word);
		}
		
		// initialize vocabulary and list of words
		InitializeVocabulary(content);
		
		// pre-calculation of all words with each count
		InitializeCounts();
		
	}
	
	// add all words to the corpus
	private void InitializeVocabulary(List<String> content)
	{
		NGram event=null;
		
		if (this.nGram==1) // handle first word
		{
			event = new NGram(content.get(0), null);
			this.words.add(event);
			this.vocabulary.put(event, 1);
			this.vocabulary_unigram.put(content.get(0), 1);
		}
		this.vocabulary_unigram.put(content.get(0), 1);
		
		for (int i=1; i<content.size(); i++) // all other words (skip index 0)
		{
			if (nGram==1)//unigram
			{
				event = new NGram(content.get(i), null);
			}
			else//bigram
			{			
				event = new NGram(content.get(i),content.get(i-1));
			}
			this.words.add(event);
			
			// add to vocabulary
			Integer count = this.vocabulary.get(event);
			Integer count_word = this.vocabulary_unigram.get(content.get(i));
			if (count == null) count = 0;
			count++;
			if (count_word == null) count_word = 0;
			count_word++;
			
			this.vocabulary.put(event, count);
			this.vocabulary_unigram.put(content.get(i), count_word);
		}
	}
	
	// in order to improve performance - pre-calculation of all words with each count
	private void InitializeCounts()
	{
		NGram event = null;
		count_r = new TreeMap<Integer,TreeSet<NGram>>();
		
		for (Entry<NGram, Integer> entry : this.vocabulary.entrySet())
		{	
			event = entry.getKey();
			int r = entry.getValue();
			
			// add r value if not exists
			if (!count_r.containsKey(r))
				count_r.put(r, new TreeSet<NGram>());
			
			count_r.get(r).add(event);
		}
	}
	
	// return the size of the corpus
	public int getSize() 
	{
		return this.words.size();
	}
	
	// counts the appearances of a distinct word in the corpus
	public int num_appearances (String word) {
		if (this.vocabulary_unigram.containsKey(word))
			return this.vocabulary_unigram.get(word);
		else
			return 0;
	}
	
	// n-gram
	public int getNGram() {
		return nGram;
	}
	
	// vocabulary size (as provided in the exercise)
	public int getVocabularySize() {
		return 300000; 
	}
	
	// return the number of distinct events appear in the text
	public int getCorpusVocabularySize() {
		return vocabulary.size();		
	}
	
	// check if an event is contained in the corpus
	public boolean contain_word(NGram event) {
		return vocabulary.containsKey(event);
	}
	
	// get distinct event keys
	public Set<NGram> getVocubularyKeySet() {
		return vocabulary.keySet();
	}
	
	// get distinct events
	public SortedMap<NGram, Integer> getVocubulary() {
		return vocabulary;
	}
	
	// get distinct words
	public Set<String> getUnigramVocubularyKeySet() {
		return vocabulary_unigram.keySet();
	}
		
	// number of appearances for an n-gram event
	public int getCount(NGram event)
	{
		Integer count = vocabulary.get(event);
		if (count == null) 
				return 0;
		
		return count.intValue();
	}
	
	// get all corpus words 
	public ArrayList<NGram> getWords() {
		return words;
	}
	
	// return list of all the events that appears r times in the corpus
	public TreeSet<NGram> words_with_count_r(int r) {
		return count_r.get(r);
	}

}
